import sys
from PyQt5.QtCore import Qt, QThread, pyqtSignal, pyqtSlot, QMutex, QWaitCondition
from PyQt5.QtWidgets import QApplication, QMainWindow, QLabel, QVBoxLayout, QWidget, QPushButton, QHBoxLayout, \
    QMessageBox, QFileDialog
from PyQt5.QtGui import QImage, QPixmap
import cv2
from ultralytics import YOLO


class Worker(QThread):
    detection_finished = pyqtSignal(object)
    video_frame = pyqtSignal(QImage, QImage)

    def __init__(self, image_path=None, video_path=None, model=None):
        super().__init__()
        self.image_path = image_path
        self.video_path = video_path
        self.model = model
        self.running = False
        self.paused = False
        self.mutex = QMutex()
        self.condition = QWaitCondition()

    def run(self):
        if self.image_path:
            self.detect_image()
        elif self.video_path:
            self.detect_video()

    def detect_image(self):
        if self.model:
            print(f"Reading image from: {self.image_path}")
            image = cv2.imread(self.image_path)
            if image is not None:
                print("Image read successfully, starting prediction...")
                try:
                    results = self.model.predict(image)
                    print("Prediction finished.")
                    self.detection_finished.emit((image, results))
                except Exception as e:
                    print(f"Prediction error: {e}")
                    self.detection_finished.emit((image, None))
            else:
                print("Failed to read image.")
                self.detection_finished.emit((None, None))

    def detect_video(self):
        if self.model:
            print(f"Reading video from: {self.video_path}")
            cap = cv2.VideoCapture(self.video_path)
            if cap.isOpened():
                self.running = True
                fps = cap.get(cv2.CAP_PROP_FPS)
                frame_delay = int(1000 / fps)  # calculate delay in milliseconds
                while self.running:
                    ret, frame = cap.read()
                    if not ret:
                        break
                    self.mutex.lock()
                    if self.paused:
                        self.condition.wait(self.mutex)
                    self.mutex.unlock()
                    try:
                        results = self.model.predict(frame)
                        annotated_frame = results[0].plot()
                        height, width, channel = frame.shape
                        bytesPerLine = 3 * width
                        qimage_original = QImage(frame.data, width, height, bytesPerLine, QImage.Format_BGR888)
                        qimage_annotated = QImage(annotated_frame.data, width, height, bytesPerLine, QImage.Format_BGR888)
                        self.video_frame.emit(qimage_original, qimage_annotated)
                        QThread.msleep(frame_delay)  # add delay between frames
                    except Exception as e:
                        print(f"Prediction error: {e}")
                        self.running = False
                cap.release()
            else:
                print("Failed to open video.")
                self.detection_finished.emit(None)

    def pause(self):
        self.paused = True

    def resume(self):
        self.paused = False
        self.condition.wakeAll()

    def stop(self):
        self.running = False
        self.resume()


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("动物模型检测")
        self.setGeometry(100, 100, 1200, 600)

        self.label_original = QLabel()
        self.label_original.setAlignment(Qt.AlignCenter)
        self.label_annotated = QLabel()
        self.label_annotated.setAlignment(Qt.AlignCenter)

        layout = QHBoxLayout()
        layout.addWidget(self.label_original)
        layout.addWidget(self.label_annotated)

        central_widget = QWidget()
        central_widget.setLayout(layout)
        self.setCentralWidget(central_widget)

        self.worker_thread = None
        self.model = None

        vbox_buttons = QVBoxLayout()
        layout.addLayout(vbox_buttons)

        self.load_model_button = QPushButton("模型选择")
        self.load_model_button.clicked.connect(self.load_model)
        vbox_buttons.addWidget(self.load_model_button)

        self.image_detect_button = QPushButton("图片检测")
        self.image_detect_button.clicked.connect(self.detect_image)
        self.image_detect_button.setEnabled(False)
        vbox_buttons.addWidget(self.image_detect_button)

        self.video_detect_button = QPushButton("视频检测")
        self.video_detect_button.clicked.connect(self.detect_video)
        self.video_detect_button.setEnabled(False)
        vbox_buttons.addWidget(self.video_detect_button)

        self.pause_button = QPushButton("暂停")
        self.pause_button.clicked.connect(self.pause_video)
        self.pause_button.setEnabled(False)
        vbox_buttons.addWidget(self.pause_button)

        self.resume_button = QPushButton("继续")
        self.resume_button.clicked.connect(self.resume_video)
        self.resume_button.setEnabled(False)
        vbox_buttons.addWidget(self.resume_button)

        self.exit_button = QPushButton("退出")
        self.exit_button.clicked.connect(self.exit_application)
        vbox_buttons.addWidget(self.exit_button)

    def detect_image(self):
        image_path, _ = QFileDialog.getOpenFileName(None, "选择图片文件", "", "图片文件 (*.jpg *.jpeg *.png)")
        if image_path:
            self.worker_thread = Worker(image_path=image_path, model=self.model)
            self.worker_thread.detection_finished.connect(self.display_image)
            self.worker_thread.start()
        else:
            QMessageBox.critical(self, "错误", "请选择图片文件")

    def detect_video(self):
        video_path, _ = QFileDialog.getOpenFileName(None, "选择视频文件", "", "视频文件 (*.mp4 *.avi)")
        if video_path:
            self.worker_thread = Worker(video_path=video_path, model=self.model)
            self.worker_thread.video_frame.connect(self.display_video_frame)
            self.worker_thread.start()
            self.pause_button.setEnabled(True)
            self.resume_button.setEnabled(True)
        else:
            QMessageBox.critical(self, "错误", "请选择视频文件")

    def display_image(self, data):
        image, results = data
        if image is not None:
            height, width, channel = image.shape
            bytesPerLine = 3 * width
            qimage_original = QImage(image.data, width, height, bytesPerLine, QImage.Format_BGR888)
            pixmap_original = QPixmap.fromImage(qimage_original)
            self.label_original.setPixmap(pixmap_original.scaled(self.label_original.size(), Qt.KeepAspectRatio))

            if results:
                print("Processing results to display image.")
                annotated_image = results[0].plot()
                qimage_annotated = QImage(annotated_image.data, width, height, bytesPerLine, QImage.Format_BGR888)
                pixmap_annotated = QPixmap.fromImage(qimage_annotated)
                self.label_annotated.setPixmap(pixmap_annotated.scaled(self.label_annotated.size(), Qt.KeepAspectRatio))
            else:
                QMessageBox.critical(self, "错误", "图像处理失败")
        else:
            QMessageBox.critical(self, "错误", "无法读取图像")

    def display_video_frame(self, qimage_original, qimage_annotated):
        pixmap_original = QPixmap.fromImage(qimage_original)
        pixmap_annotated = QPixmap.fromImage(qimage_annotated)
        self.label_original.setPixmap(pixmap_original.scaled(self.label_original.size(), Qt.KeepAspectRatio))
        self.label_annotated.setPixmap(pixmap_annotated.scaled(self.label_annotated.size(), Qt.KeepAspectRatio))

    def pause_video(self):
        if self.worker_thread:
            self.worker_thread.pause()

    def resume_video(self):
        if self.worker_thread:
            self.worker_thread.resume()

    def load_model(self):
        model_path, _ = QFileDialog.getOpenFileName(None, "选择模型文件", "", "模型文件 (*.pt)")
        if model_path:
            print(f"Loading model from: {model_path}")
            try:
                self.model = YOLO(model_path)
                if self.model:
                    print("Model loaded successfully.")
                    self.image_detect_button.setEnabled(True)
                    self.video_detect_button.setEnabled(True)
                else:
                    print("Model loading failed.")
                    QMessageBox.critical(self, "错误", "模型加载失败")
            except Exception as e:
                print(f"Model loading error: {e}")
                QMessageBox.critical(self, "错误", f"模型加载失败: {e}")

    def exit_application(self):
        if self.worker_thread is not None:
            self.worker_thread.stop()
        sys.exit()

    def resizeEvent(self, event):
        if self.label_original.pixmap():
            self.label_original.setPixmap(
                self.label_original.pixmap().scaled(self.label_original.size(), Qt.KeepAspectRatio))
        if self.label_annotated.pixmap():
            self.label_annotated.setPixmap(
                self.label_annotated.pixmap().scaled(self.label_annotated.size(), Qt.KeepAspectRatio))
        super().resizeEvent(event)


if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())
